from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.status import HTTP_500_INTERNAL_SERVER_ERROR
import time
import logging
from typing import Awaitable, Callable

# 获取日志器
logger = logging.getLogger("fastapi_app")

class LoggingMiddleware(BaseHTTPMiddleware):
    # 响应内容的最大记录长度
    MAX_RESPONSE_BODY_LENGTH = 5000  # 5KB

    async def dispatch(self, request: Request, call_next: Callable[[Request], Awaitable[Response]]) -> Response:
        # 记录请求开始时间
        start_time = time.time()

        # 记录请求信息
        logger.info(
            f"请求: 方法={request.method}, 路径={request.url.path}, "
            f"客户端IP={request.client.host}, 查询参数={dict(request.query_params)}"
        )

        try:
            # 处理请求
            response = await call_next(request)

            # 计算请求处理时间
            process_time = time.time() - start_time

            # 读取响应内容
            # response_body = await self._get_response_body(response)

            # 记录响应信息
            logger.info(
                f"响应: 状态码={response.status_code}, 处理时间={process_time:.6f}秒, "
                f"方法={request.method}, 路径={request.url.path}"
                # f"方法={request.method}, 路径={request.url.path}, 内容={response_body}"
            )

            return response

        except Exception as e:
            # 记录异常信息
            process_time = time.time() - start_time
            logger.error(
                f"请求处理出错: 方法={request.method}, 路径={request.url.path}, "
                f"错误={str(e)}, 处理时间={process_time:.6f}秒",
                exc_info=True  # 记录详细的异常堆栈信息
            )

            # 返回500错误响应
            return Response(
                content="服务器内部错误",
                status_code=HTTP_500_INTERNAL_SERVER_ERROR
            )

    async def _get_response_body(self, response: Response) -> str:
        """获取响应内容并进行适当处理"""
        # 复制响应体以便我们可以读取它
        body = b""
        async for chunk in response.body_iterator:
            body += chunk

        # 将内容转换回迭代器
        from starlette.responses import StreamingResponse
        response.body_iterator = StreamingResponse(iter([body])).body_iterator

        # 转换为字符串并截断过长的内容
        body_str = body.decode()
        if len(body_str) > self.MAX_RESPONSE_BODY_LENGTH:
            return f"{body_str[:self.MAX_RESPONSE_BODY_LENGTH]}... [内容过长，已截断]"

        return body_str
